package xyz.naokeziteng.data_structure.heap;

import java.util.ArrayList;
import java.util.Arrays;

/**
 * @author hawk
 * @date 2022/9/7
 * @desc 二叉堆
 **/
public class MaxHeap<E extends Comparable<E>> {
    //二叉堆是一棵完全二叉树
    //层序遍历方式存放，缺失的一定是右子树
    //数组存放
    //parent(i) = (i-1)/2
    //left child(i) = 2*i + 1
    //right child(i) = 2*i + 2

    private ArrayList<E> data;

    public MaxHeap(int capacity) {
        data = new ArrayList<>(capacity);
    }

    public MaxHeap() {
        data = new ArrayList<>();
    }

    public MaxHeap(E[] data) {
        this.data = new ArrayList<>(Arrays.asList(data));
        heapify();
    }

    private void heapify() {
        int nonLeafIndex = ((data.size() - 1) - 1) / 2;
        while (nonLeafIndex >= 0) {
            siftDown(nonLeafIndex);
            nonLeafIndex--;
        }
    }

    public int size() {
        return data.size();
    }

    public boolean isEmpty() {
        return data.isEmpty();
    }

    public E findMax() {
        if (size() == 0) {
            throw new IllegalArgumentException("empty heap do not have max leaf...");
        }
        return data.get(0);
    }

    //返回完全二叉树的数组表示中，一个索引所表示的元素的父亲节点的索引
    private int parent(int index) {
        if (index == 0) {
            throw new IllegalArgumentException("根节点没有父节点...");
        }
        return (index - 1) / 2;
    }

    //返回完全二叉树的数组表示中，一个索引所表示的元素的左孩子节点的索引
    private int leftChild(int index) {
        return index * 2 + 1;
    }

    //返回完全二叉树的数组表示中，一个索引所表示的元素的右孩子节点的索引
    private int rightChild(int index) {
        return index * 2 + 2;
    }

    public void add(E e) {
        data.add(e);
        siftUp(data.size() - 1);
    }

    public E poll() {
        if (size() == 0) {
            throw new IllegalArgumentException("empty heap can not poll...");
        }
        E result = data.get(0);
        swap(0, size() - 1);
        data.remove(size() - 1);
        siftDown(0);
        return result;
    }

    //取出最大元素，并替换成新的元素
    public E replace(E e) {
        if (size() == 0) {
            throw new IllegalArgumentException("empty heap can not poll...");
        }
        E result = data.get(0);
        data.set(0, e);
        siftDown(0);
        return result;
    }


    private void siftUp(int index) {
        while (index > 0) {
            if (data.get(parent(index)).compareTo(data.get(index)) < 0) {
                //交换
                swap(parent(index), index);
                index = parent(index);
                continue;
            }
            break;
        }
    }

    private void swap(int i, int j) {
        E tmp = data.get(i);
        data.set(i, data.get(j));
        data.set(j, tmp);
    }

    private void siftDown(int index) {
        while (index * 2 + 1 < size()) {
            int j = index * 2 + 1;
            if (j + 1 < size() && data.get(j).compareTo(data.get(j + 1)) < 0) {
                j++;
            }
            if (data.get(index).compareTo(data.get(j)) < 0) {
                swap(index, j);
                index = j;
                continue;
            }
            break;
        }
    }

    public static void main(String[] args) {
        MaxHeap<Integer> maxHeap = new MaxHeap<>();
        maxHeap.add(9);
        maxHeap.add(5);
        maxHeap.add(8);
        maxHeap.add(4);
        maxHeap.add(2);
        System.out.println(maxHeap.data);
        maxHeap.add(10);
        System.out.println(maxHeap.data);
        System.out.println(maxHeap.poll());
        System.out.println(maxHeap.data);

    }

}
